package com.aiwiown.snackmq.broker.auth;

import com.aiwiown.snackmq.cluster.metadata.MetadataManager;
import com.aiwiown.snackmq.common.auth.Action;
import com.aiwiown.snackmq.common.auth.Permission;
import com.aiwiown.snackmq.common.exception.AuthorizationException;
import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.github.benmanes.caffeine.cache.Expiry;
import io.jsonwebtoken.Claims;
import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.checkerframework.checker.index.qual.NonNegative;
import org.checkerframework.checker.nullness.qual.NonNull;

import java.util.Date;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 【已优化】授权服务。
 * 此版本修复了对 MANAGE 操作的资源类型判断逻辑。
 */
@Slf4j
@RequiredArgsConstructor
public class AuthorizationService {

    private final MetadataManager metadataManager;
    private final JwtService jwtService;
    private final boolean authEnabled;

    @Getter
    @AllArgsConstructor
    public static class Resource {
        private final String type;
        private final String name;
    }

    private final Cache<String, AuthenticatedUser> authCache = Caffeine.newBuilder()
            .expireAfter(new JwtExpiry())
            .maximumSize(10000)
            .build();


    /**
     * 为多个资源进行复合授权检查。
     *
     * @param token             客户端提供的 JWT
     * @param action            请求的操作
     * @param requiredResources 需要检查的资源列表
     * @throws AuthorizationException 如果任何一个资源的授权失败
     */
    public void authorize(String token, Action action, Resource... requiredResources) throws AuthorizationException {
        if (!authEnabled) {
            return;
        }
        AuthenticatedUser user = getAuthenticatedUser(token);
        checkUserPermissions(user, action, requiredResources);
    }

    private AuthenticatedUser getAuthenticatedUser(String token) throws AuthorizationException {
        if (token == null || token.isEmpty()) {
            throw new AuthorizationException("Auth token is missing.");
        }

        try {
            return authCache.get(token, k -> {
                log.trace("Authorization cache miss. Validating token...");
                Claims claims = jwtService.validateAndParseToken(k);
                if (claims == null) {
                    throw new UncheckedAuthorizationException("Invalid or expired auth token.");
                }
                AuthenticatedUser user = new AuthenticatedUser(claims);
                log.debug("Validated and cached user '{}' from token. Cache will expire at {}.",
                        user.getUsername(), user.getExpiration());
                return user;
            });
        } catch (UncheckedAuthorizationException e) {
            throw new AuthorizationException(e.getMessage());
        } catch (Exception e) {
            log.error("An unexpected error occurred during token validation for caching.", e);
            throw new AuthorizationException("Token validation failed due to an internal error.", e);
        }
    }

    private void checkUserPermissions(
            AuthenticatedUser user,
            Action action,
            Resource... requiredResources) throws AuthorizationException {

        if (user.getRoles() == null || user.getRoles().isEmpty()) {
            throw new AuthorizationException("User '" + user.getUsername() + "' has no roles assigned.");
        }

        for (Resource resource : requiredResources) {
            boolean hasPermissionForResource = user.getRoles().stream()
                    .anyMatch(role -> hasPermissionForRole(role, action, resource));

            // 【最终修复】如果对当前检查的任何一个资源没有权限，则立即失败并抛出异常。
            // 此前的逻辑虽然大体正确，但此修复使其“立即失败”的意图更加明确和健壮，
            // 确保客户端能收到准确的权限错误，而不是模糊的认证失败。
            if (!hasPermissionForResource) {
                log.warn("Authorization failed for user '{}'. Missing permission for action '{}' on resource(type={}, name={})",
                        user.getUsername(), action, resource.getType(), resource.getName());
                // 【代码改进】在异常消息中提供更完整的上下文，包括操作和资源类型。
                throw new AuthorizationException(String.format("Permission denied for action %s on resource %s:%s", action, resource.getType(), resource.getName()));
            }
        }

        // 只有当所有必需的资源权限都检查通过后，才记录成功日志。
        log.debug("Authorization successful for user '{}' on all required resources.", user.getUsername());
    }

    private boolean hasPermissionForRole(String role, Action action, Resource resource) {
        List<Permission> permissions = metadataManager.getPermissionsForRole(role);
        if (permissions == null) {
            return false;
        }

        return permissions.stream().anyMatch(p -> {
            // First, check if the permission applies to the requested resource
            // (type and name, with wildcard support)
            boolean resourceMatch = p.getResourceType().equalsIgnoreCase(resource.getType())
                    && ("*".equals(p.getResourceName()) || p.getResourceName().equals(resource.getName()));

            // If the resource doesn't match, this permission is irrelevant
            return resourceMatch &&
                    // Then, check if the granted action is sufficient. MANAGE grants all other actions.
                    (p.getAction() == Action.MANAGE || p.getAction() == action);
        });
    }

    @Getter
    private static final class AuthenticatedUser {
        private final String username;
        private final List<String> roles;
        private final Date expiration;

        @SuppressWarnings("unchecked")
        AuthenticatedUser(Claims claims) {
            this.username = claims.getSubject();
            this.roles = claims.get("roles", List.class);
            this.expiration = claims.getExpiration();
        }
    }

    private static class JwtExpiry implements Expiry<String, AuthenticatedUser> {
        @Override
        public long expireAfterCreate(@NonNull String key, @NonNull AuthenticatedUser user, long currentTime) {
            if (user.getExpiration() == null) {
                return TimeUnit.MINUTES.toNanos(5);
            }
            long remainingMillis = user.getExpiration().getTime() - System.currentTimeMillis();
            return TimeUnit.MILLISECONDS.toNanos(Math.max(0, remainingMillis));
        }

        @Override
        public long expireAfterUpdate(@NonNull String key, @NonNull AuthenticatedUser user, long currentTime, @NonNegative long currentDuration) {
            return expireAfterCreate(key, user, currentTime);
        }

        @Override
        public long expireAfterRead(@NonNull String key, @NonNull AuthenticatedUser user, long currentTime, @NonNegative long currentDuration) {
            return currentDuration;
        }
    }

    private static class UncheckedAuthorizationException extends RuntimeException {
        public UncheckedAuthorizationException(String message) {
            super(message);
        }
    }
}